import pySALEPlot as psp
from pylab import figure,arange,colorbar,loadtxt,figtext
# This allows the figure to be subdivided nonuniformly
from matplotlib import gridspec
# Need this for the colorbars we will make on the mirrored plot
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Make an output directory
dirname='Plots'
psp.mkdir_p(dirname)

# Specify the set of simulations to process
simulations = ["Dilatancy"]
# Plot title (should be same length as "simualtions")
ptitle = ["$L$ = 11.6  km; $v_i$ = 13 km/s; $g$ = 3.71 m/s$^2$"]
# x-limit of the plot (should be same length as "simualtions")
xmax = [32.]

n=0
for sim in simulations:

    # Define datafilename
    datafilename = 'jdata.dat'

    # Open the datafile
    model=psp.opendatfile(datafilename)

    # Set the distance units to km
    model.setScale('km')

    # If a single image with no gravity profile, create a simple single axes figure.
    fig=figure(figsize=(8,6))
    gs = gridspec.GridSpec(1,1)
    ax=fig.add_subplot(gs[0],aspect='equal')

    for i in arange(0,model.nsteps,1):

        # Set the axis labels
        ax.set_xlabel('Distance from symmetry axis (km)')
        ax.set_ylabel('Height (km)')

	ax.set_title(ptitle[n])

        # Set the axis limits
        ax.set_xlim([0,90])
        ax.set_ylim([-40,20])

        # Read the step
        step=model.readStep('Alp',i)
	print 'Processing step: ',i

        # Plot the distension field as porosity as a colourmap
        p=ax.pcolormesh(model.x,model.y,1.-1./step.data[0],
            cmap='binary',vmin=0.,vmax=0.2)

        # And the material boundary
        b=ax.contour(model.xc,model.yc,step.cmc[0],colors='k',levels=[0.1])
        c=ax.contour(model.xc,model.yc,step.cmc[1],colors='k',levels=[0.1])
   
	# Add a colorbar (only need to do this once)
	if (i == 0):
            cb=fig.colorbar(p,orientation='horizontal',shrink=0.8,pad=0.15)
            cb.set_label('Porosity')

        # Add a tracer overlay
        for u in range(model.tracer_numu):
            tru=model.tru[u]
            # Plot the tracers in horizontal lines, every 10 lines
            for l in arange(0,len(tru.xlines),10):
               ax.plot(step.xmark[tru.xlines[l]],
                       step.ymark[tru.xlines[l]],
                       c='k',marker=',',linestyle='None',markersize=0.5)
            # Plot the tracers in vertical lines, every 10 lines
            for l in arange(0,len(tru.ylines),10):
                ax.plot(step.xmark[tru.ylines[l]],
                        step.ymark[tru.ylines[l]],
                        c='k',marker=',',linestyle='None',markersize=0.5)

        # Add time labels to the frames
        ax.annotate('T = {: 3.0f} s'.format(step.time),xy=(0.775*xmax[n],xmax[n]/6.),xycoords='data',horizontalalignment='left',verticalalignment='top')

        # Save the figure if one per simulation
        fig.savefig('./{}/{}-{:05d}.png'.format(dirname,sim,i),dpi=300)

	# Remove the field, ready for the next timestep to be plotted
        ax.cla()

    n = n+1
